from __future__ import annotations

from finrl.meta.env_stock_trading.env_stock_papertrading import AlpacaPaperTrading
from finrl.test import test
import numpy as np
import baostock as bs
from datetime import datetime, timedelta
import time

class THSPaperTrading:
    def __init__(
            self,
            ticker_list,
            time_interval,
            drl_lib,
            model_name,
            cwd,
            net_dim,
            state_dim,
            action_dim,
            initial_amount=1000000,  # 初始资金
            max_stock=100000,        # 最大持仓数量
            latency=0.1             # 交易延迟
    ):
        """
        初始化同花顺模拟交易环境

        Args:
            ticker_list (list): 股票代码列表
            time_interval (str): 交易时间间隔，如 '1d', '1h', '5m'
            drl_lib (str): 深度强化学习库名称
            model_name (str): 模型名称
            cwd (str): 当前工作目录
            net_dim (int): 神经网络维度
            state_dim (int): 状态空间维度
            action_dim (int): 动作空间维度
            initial_amount (float): 初始资金
            max_stock (int): 最大持仓数量
            latency (float): 交易延迟（秒）
        """
        # 基础配置
        self.ticker_list = ticker_list
        self.time_interval = time_interval
        self.initial_amount = initial_amount
        self.max_stock = max_stock
        self.latency = latency

        # 账户状态
        self.stocks = np.asarray([0] * len(ticker_list))  # 持仓量
        self.stocks_cd = np.zeros_like(self.stocks)       # 交易冷却时间（T+1）
        self.cash = initial_amount                        # 现金余额

        # 加载模型
        if drl_lib == 'stable_baselines3':
            from stable_baselines3 import A2C, DDPG, PPO, SAC, TD3

            if model_name == 'a2c':
                self.model = A2C.load(cwd)
            elif model_name == 'ddpg':
                self.model = DDPG.load(cwd)
            elif model_name == 'td3':
                self.model = TD3.load(cwd)
            elif model_name == 'sac':
                self.model = SAC.load(cwd)
            elif model_name == 'ppo':
                self.model = PPO.load(cwd)
            else:
                raise ValueError(f'未知模型: {model_name}')
        else:
            raise ValueError(f'未支持的DRL库: {drl_lib}')

        # 市场数据源
        self.data_source = self._init_data_source()

    def _init_data_source(self):
        """初始化数据源"""
        bs.login()
        return bs

    def get_state(self):
        """获取当前市场状态"""
        state = np.zeros(len(self.ticker_list) * 3 + 1)

        # 现金比例
        state[0] = self.cash / self.initial_amount

        for i, ticker in enumerate(self.ticker_list):
            # 获取实时价格
            current_price = self.get_real_time_price(ticker)
            position = self.stocks[i]

            # 更新状态向量
            state[3 * i + 1] = position / self.max_stock
            state[3 * i + 2] = current_price
            state[3 * i + 3] = self.stocks_cd[i]

        return state

    def get_real_time_price(self, ticker):
        """获取实时价格"""
        try:
            # 转换股票代码格式
            if ticker.endswith('.SS'):
                bs_code = f"sh.{ticker[:-3]}"
            elif ticker.endswith('.SZ'):
                bs_code = f"sz.{ticker[:-3]}"

            rs = self.data_source.query_latest_trade_data(code=bs_code)
            if rs.error_code == '0':
                return float(rs.get_row_data()[5])  # 收盘价
            else:
                print(f"获取{ticker}实时价格失败")
                return None
        except Exception as e:
            print(f"获取实时价格出错: {e}")
            return None

    def calculate_trade_cost(self, price, quantity, is_buy=True):
        """计算交易成本"""
        total = price * quantity

        # 佣金（双向收取）
        commission = max(total * 0.00025, 5)

        # 印花税（卖出时收取）
        stamp_duty = total * 0.001 if not is_buy else 0

        # 过户费（上海市场）
        transfer_fee = total * 0.00002 if ticker.startswith('6') else 0

        return commission + stamp_duty + transfer_fee

    def execute_trades(self, actions):
        """执行交易"""
        for i, action in enumerate(actions):
            ticker = self.ticker_list[i]
            current_price = self.get_real_time_price(ticker)

            if current_price is None:
                continue

            # 计算目标持仓
            target_amount = self.max_stock * action
            delta_amount = target_amount - self.stocks[i]

            # T+1 限制检查
            if delta_amount < 0 and self.stocks_cd[i] > 0:
                print(f"{ticker} T+1限制，无法卖出")
                continue

            # 涨跌停检查
            if self.is_price_limit(ticker, current_price):
                print(f"{ticker} 涨跌停限制")
                continue

            # 执行交易
            if delta_amount > 0:  # 买入
                cost = self.calculate_trade_cost(current_price, delta_amount, True)
                total_cost = current_price * delta_amount + cost

                if total_cost <= self.cash:
                    self.cash -= total_cost
                    self.stocks[i] += delta_amount
                    self.stocks_cd[i] = 1  # 设置T+1冷却
                    print(f"买入 {ticker}: {delta_amount}股，价格：{current_price}")

            elif delta_amount < 0:  # 卖出
                delta_amount = abs(delta_amount)
                cost = self.calculate_trade_cost(current_price, delta_amount, False)
                total_revenue = current_price * delta_amount - cost

                self.cash += total_revenue
                self.stocks[i] -= delta_amount
                print(f"卖出 {ticker}: {delta_amount}股，价格：{current_price}")

    def is_price_limit(self, ticker, current_price):
        """检查是否涨跌停"""
        # 获取昨日收盘价
        try:
            if ticker.endswith('.SS'):
                bs_code = f"sh.{ticker[:-3]}"
            elif ticker.endswith('.SZ'):
                bs_code = f"sz.{ticker[:-3]}"

            rs = self.data_source.query_history_k_data_plus(
                code=bs_code,
                fields="close",
                start_date=(datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d'),
                end_date=datetime.now().strftime('%Y-%m-%d')
            )

            if rs.error_code == '0' and rs.next():
                prev_close = float(rs.get_row_data()[0])
                # 涨跌停限制（一般是10%，科创板20%）
                limit = 0.2 if ticker.startswith('688') else 0.1
                upper_limit = prev_close * (1 + limit)
                lower_limit = prev_close * (1 - limit)

                return current_price >= upper_limit or current_price <= lower_limit

        except Exception as e:
            print(f"检查涨跌停出错: {e}")
            return True  # 出错时保守处理，认为是涨跌停

        return False

    def run(self):
        """运行交易系统"""
        try:
            while True:
                # 检查是否是交易时间
                if not self.is_market_open():
                    print("市场休市中...")
                    time.sleep(60)
                    continue

                # 获取当前状态
                state = self.get_state()

                # 模型预测
                actions = self.model.predict(state)[0]

                # 执行交易
                self.execute_trades(actions)

                # 更新账户状态
                self.update_account_status()

                # 打印当前持仓信息
                self.print_portfolio()

                # 等待下一个时间间隔
                time.sleep(self.latency)

        except KeyboardInterrupt:
            print("手动停止交易")
        finally:
            self.data_source.logout()

    def is_market_open(self):
        """检查市场是否开市"""
        now = datetime.now()
        weekday = now.weekday()

        # 周末休市
        if weekday >= 5:
            return False

        # 交易时间：9:30-11:30, 13:00-15:00
        time = now.time()
        morning_start = time(9, 30)
        morning_end = time(11, 30)
        afternoon_start = time(13, 0)
        afternoon_end = time(15, 0)

        return (morning_start <= time <= morning_end) or \
            (afternoon_start <= time <= afternoon_end)

    def update_account_status(self):
        """更新账户状态"""
        # 更新T+1冷却时间
        self.stocks_cd = np.maximum(0, self.stocks_cd - 1)

        # 计算总资产
        total_value = self.cash
        for i, ticker in enumerate(self.ticker_list):
            price = self.get_real_time_price(ticker)
            if price is not None:
                total_value += price * self.stocks[i]

        return total_value

    def print_portfolio(self):
        """打印投资组合信息"""
        print("\n=== 当前持仓信息 ===")
        print(f"现金: {self.cash:.2f}")

        for i, ticker in enumerate(self.ticker_list):
            if self.stocks[i] > 0:
                price = self.get_real_time_price(ticker)
                value = price * self.stocks[i] if price else 0
                print(f"{ticker}: {self.stocks[i]}股, 现值: {value:.2f}")

        total_value = self.update_account_status()
        print(f"总资产: {total_value:.2f}")
        print("==================\n")
